package com.yahoo.labs.snow.util;

import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
import it.unimi.dsi.io.FileLinesCollection.FileLinesIterator;
import it.unimi.dsi.lang.MutableString;
import it.unimi.dsi.logging.ProgressLogger;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintWriter;

import org.apache.log4j.Logger;

import com.martiansoftware.jsap.FlaggedOption;
import com.martiansoftware.jsap.JSAP;
import com.martiansoftware.jsap.JSAPException;
import com.martiansoftware.jsap.JSAPResult;
import com.martiansoftware.jsap.Parameter;
import com.martiansoftware.jsap.SimpleJSAP;
import com.martiansoftware.jsap.Switch;

/**
 * Computer Jaccard similarity.
 * 
 * It takes an input file containing pairs of the form <tt>(key,value)</tt>.
 * 
 * It outputs pairs of <tt>(key1,key2,jaccard(values(key1),values(key2))</tt>, using the jaccard
 * similarity (size of intersection divided by size of union), where <tt>values(key)</tt> are all
 * the distinct values seen for the given key in the input.
 * 
 * See --help for options.
 * 
 * @author chato
 * 
 */
public class PairsToJaccardSimilarity {

	static Logger logger = Logger.getLogger(PairsToJaccardSimilarity.class);

	final static String DEFAULT_FIELD_SEPARATOR = "\t";

	public static void main(String[] args) throws JSAPException, FileNotFoundException {
		final SimpleJSAP jsap = new SimpleJSAP(PairsToJaccardSimilarity.class.getName(), "Process pairs of items .", new Parameter[] {
				new FlaggedOption("input", JSAP.STRING_PARSER, JSAP.NO_DEFAULT, JSAP.REQUIRED, 'i', "input", "Input file of pairs <key,value>."),
				new FlaggedOption("output", JSAP.STRING_PARSER, JSAP.NO_DEFAULT, JSAP.REQUIRED, 'o', "output",
						"Output file of pairs <key1,key2,jaccard(key1.value,key2.value)>."),
				new Switch("intersection", JSAP.NO_SHORTFLAG, "intersection", "Computes the size of intersection instead of the Jaccard coefficient"),
				new FlaggedOption("key", JSAP.INTEGER_PARSER, Integer.toString(0), JSAP.NOT_REQUIRED, 'k', "key", "Number of input field corresponding to key, start from 0."),
				new FlaggedOption("field-separator", JSAP.STRING_PARSER, DEFAULT_FIELD_SEPARATOR, JSAP.NOT_REQUIRED, 't', "field-separator", "Field separator."), });

		final JSAPResult params = jsap.parse(args);
		if (jsap.messagePrinted()) {
			return;
		}

		File inFile = new File(params.getString("input"));
		File outFile = new File(params.getString("output"));
		String fieldSeparator = params.getString("field-separator");
		int keyField = params.getInt("key");
		if (keyField < 0 || keyField > 1) {
			throw new IllegalArgumentException("The key field must be 0 or 1");
		}
		int valueField = keyField == 0 ? 1 : 0;
		Object2ObjectOpenHashMap<String, ObjectOpenHashSet<String>> sets = new Object2ObjectOpenHashMap<String, ObjectOpenHashSet<String>>();

		logger.info("Key field=" + keyField + ", value field=" + valueField);
		
		boolean intersection = false;
		if( params.userSpecified("intersection") ) {
			intersection = params.getBoolean("intersection");
			if( intersection ) {
				logger.info("Will compute size of intersection instead of Jaccard coefficient");
			}
		}

		logger.info("Reading input from " + inFile);
		FileLinesIterator it = FileUtils.getIterator(inFile);
		while (it.hasNext()) {
			MutableString line = it.next();
			String[] tokens = line.toString().split(fieldSeparator);
			String key = tokens[keyField];
			String value = tokens[valueField];
			if (!sets.containsKey(key)) {
				sets.put(key, new ObjectOpenHashSet<String>());
			}
			sets.get(key).add(value);
		}

		logger.info("Computing similarities and writing them to " + outFile);
		String[] keys = (sets.keySet()).toArray(new String[] {});
		ProgressLogger pl = new ProgressLogger(logger, ProgressLogger.TEN_SECONDS, "nodes");
		pl.expectedUpdates = keys.length;
		pl.start();
		PrintWriter pw = FileUtils.getWriter(outFile);
		for (int i=0; i<keys.length; i++ ) {
			pl.update();
			String k1 = keys[i];
			ObjectOpenHashSet<String> s1 = sets.get(k1);
			for( int j=i+1; j<keys.length; j++ ) {
				String k2 = keys[j];
				ObjectOpenHashSet<String> s2 = sets.get(k2);
				double value = intersection ? computeIntersection(s1,s2) : computeJaccard(s1,s2);
				if( value > 0 ) {
					pw.println(k1 + fieldSeparator + k2 + fieldSeparator + value);
				}
			}
		}
		for (String k1 : sets.keySet()) {
			pl.update();
			for (String k2 : sets.keySet()) {
				if (!k1.equals(k2)) {
				
				}
			}
		}
		pl.stop();
		pw.close();
	}
	
	static double computeIntersection(ObjectOpenHashSet<String> set1, ObjectOpenHashSet<String> set2) {
		ObjectOpenHashSet<String> inter = new ObjectOpenHashSet<String>((int) Math.min(set1.size(), set2.size()));

		ObjectOpenHashSet<String> smallerSet = set1.size() > set2.size() ? set2 : set1;
		ObjectOpenHashSet<String> largerSet = set1.size() > set2.size() ? set1 : set2;

		for (String elem : smallerSet) {
			if (largerSet.contains(elem)) {
				inter.add(elem);
			}
		}
		
		return( inter.size() );
	}
	
	static double computeJaccard(ObjectOpenHashSet<String> set1, ObjectOpenHashSet<String> set2) {
		ObjectOpenHashSet<String> union = new ObjectOpenHashSet<String>(set1.size() + set2.size());
		ObjectOpenHashSet<String> inter = new ObjectOpenHashSet<String>((int) Math.min(set1.size(), set2.size()));

		ObjectOpenHashSet<String> smallerSet = set1.size() > set2.size() ? set2 : set1;
		ObjectOpenHashSet<String> largerSet = set1.size() > set2.size() ? set1 : set2;

		for (String elem : smallerSet) {
			union.add(elem);
			if (largerSet.contains(elem)) {
				inter.add(elem);
			}
		}
		
		// If the intersection is zero the jaccard is zero
		if( inter.size() == 0 ) {
			return 0.0;
		}
		
		for (String elem : largerSet) {
			if( ! union.contains(elem) ) {
				union.add(elem);
			}
		}

		// The union is certainly larger than 0 if the intersection is larger than zero
		return ((double) inter.size() / (double) union.size());
	}

}
